package org.base23.ai.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.time.LocalDate;

@RestController
@RequestMapping("/ai")
@Tag(name = "AI服务", description = "AI聊天和生成服务")
public class AIController {

    private final ChatClient chatClient;

    public AIController(ChatClient.Builder chatClientBuilder, ChatMemory chatMemory, VectorStore vectorStore, @Value("${ai.prompt}") String prompt) {
        this.chatClient = chatClientBuilder.defaultSystem(prompt)
                .defaultAdvisors(
                        new PromptChatMemoryAdvisor(chatMemory)
                        //new LoggingAdvisor()
                        ,new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults())
                )
//                .defaultFunctions("myFunc")
                .build();
    }

    @CrossOrigin
    @GetMapping(value = "/generateStreamAsString", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @Operation(summary = "流式生成回答", description = "使用SpringAI流式生成回答")
    public Flux<String> generateStreamAsString(@RequestParam(value = "message", defaultValue = "你好") String message) {
        return this.chatClient.prompt()
                .user(message)
                .system(promptSystemSpec -> promptSystemSpec.param("current_date", LocalDate.now().toString()))
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                //.advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.query(message)))
                .stream()
                .content()
                .onErrorResume(e -> {
                    System.err.println("Error during chat: " + e.getMessage());
                    return Flux.just("抱歉，我遇到了一些问题，请稍后再试。");
                })
                .concatWith(Flux.just("[complete]"));
    }
}
